import sys
from pathlib import Path

base_path = str(Path(__file__).resolve().parent.parent)
sys.path.append(base_path)
from pydantic import BaseModel, Field
from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import (
    PromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
)
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage
from langchain_text_splitters import MarkdownHeaderTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from typing import List
from create_llm import create_embeddings, create_llm

# 要加载的两个markdown文件路径
file_path_list = [
    f"{base_path}/data/md_list/百炼_图像局部重绘.md",
    f"{base_path}/data/md_list/百炼_涂鸦作画.md",
]


def define_history():
    history_list: List[BaseMessage] = []

    def get_history_list():
        return history_list

    def save_history_item(message):
        if isinstance(message, BaseMessage):
            history_list.append(message)

    return (get_history_list, save_history_item)


get_history_list, save_history_item = define_history()


# 定义用Pydantic约束的llm返回问题模型
class RetrieveInfo(BaseModel):
    retrieve_flag: str = Field(description="是否需要检索文档，值为Y或N")
    code_flag: str = Field(description="是否为代码类的问题，值为Y或N")


# 用RetrieveInfo定义JSON解析器
retrieve_parser = JsonOutputParser(pydantic_object=RetrieveInfo)


# 加载文档并按markdown分割，返回由分割chunk组成的Document数组
def load_and_split_document():
    # 定义返回chunk数组
    documents = []
    # 定义markdown文档分割器
    splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=[
            ("#", "一级标题"),
            ("##", "二级标题"),
            ("###", "三级标题"),
        ],
        strip_headers=False,
    )
    for file_path in file_path_list:
        # 根据文件路径定义文档加载器
        loader = TextLoader(file_path=file_path, encoding="utf-8")
        # 加载文档
        file_content = loader.load()
        file_content = file_content[0].page_content
        # 分割文档
        chunks = splitter.split_text(file_content)
        # 添加进documents数组
        documents.extend(chunks)

    return documents


# 文档列表各chunk使用嵌入式模型转为向量数据并存入向量数据库
def create_vector_store_instance(documents):
    # 创建嵌入式模型实例
    embeddings = create_embeddings()
    # 基于文档列表、嵌入式模型实例创建向量数据库实例
    chroma_instance = Chroma.from_documents(documents=documents, embedding=embeddings)
    return chroma_instance


# 根据用户的提问与向量数据库中存储chunks的相关性，返回对应chunk
def search_vector(vector_instance, question, isCode=False):
    if isCode:
        # 代码类问题，使用最大边际相关性搜索算法获取相关文档，k=3表示返回三篇文档
        docs = vector_instance.max_marginal_relevance_search(query=question, k=3)
    else:
        # 非代码问题，使用相似度搜索算法获取相关文档，k=3表示返回3篇最相关文档
        docs = vector_instance.similarity_search(query=question, k=3)

    # 异常检测，若为检索到相关内容，则返回None
    if len(docs) == 0:
        return None

    # 最后筛选出的chunk组成的数组
    last_docs = []
    for doc in docs:
        last_docs.append(doc.page_content)
    # 使用5个换行符连接last_docs，组成一个字符串
    return ("\n" * 5).join(last_docs)


# 测试文档加载、文档分割、chunk向量化后入库、根据提问从向量数据库中搜索对应文档等功能
# def test():
#     document = load_and_split_document()
#     chroma_instance = create_vector_store_instance(document)
#     result = search_vector(chroma_instance, "百炼涂鸦作画的java实现", isCode=True)
#     print(result)
# test()


def log(val):
    print(val)
    print("\n")
    print("-" * 50)
    return val


flags_chain = (
    PromptTemplate.from_template(
        """
# 请根据用户的提问和文件列表中的文件名的匹配结果按要求生成JSON字符串并返回
## 用户提问：
  {question}                                   
## 文件列表：
  {file_paths}                                            
## 要求：

  JSON有retrieve_flag和code_flag两个属性。
  如果用户提问和文件名存在相同名词，则retrieve_flag属性值为Y，不存在，则retrieve_flag属性值为N。
  如果用户提问类型是代码类问题，则code_flag属性值为Y，如果不是，code_flag属性值为N。
"""
    )
    | create_llm(temperature=0)
    # | log
    | retrieve_parser
    # | log
)


# 根据用户提问和上一步对用户提问和文件名匹配的情况得到的flags调用search_vector检索出用户提问在向量数据库中相近的文档块作为参考信息给到大模型
def get_result_prompt_text(data):
    question = data["question"]
    flags = data["flags"]
    retrieve_flag = flags["retrieve_flag"]
    code_flag = flags["code_flag"]

    document = load_and_split_document()
    chroma_instance = create_vector_store_instance(document)
    chunks_text = None
    if retrieve_flag == "Y" and code_flag == "Y":
        chunks_text = search_vector(chroma_instance, question, isCode=True)
    elif retrieve_flag == "Y" and code_flag == "N":
        chunks_text = search_vector(chroma_instance, question, isCode=False)
    return f"""
    # 请根据用户的提问和给出的参考文件，整理思考后返回用户最终的答案
    ## 用户提问：
      {question}
    ## 参考文件：
      {chunks_text}
    """


prompt_chain = (
    RunnablePassthrough.assign(
        history=lambda x: get_history_list(),
        file_paths=lambda x: ",".join(file_path_list),
    )
    | RunnablePassthrough.assign(flags=flags_chain)
    # | log
    | RunnablePassthrough.assign(result_prompt_text=lambda x: get_result_prompt_text(x))
    | ChatPromptTemplate.from_messages(
        [
            SystemMessage(content="你是一名智能助手，你需要解答用户的提问"),
            MessagesPlaceholder(variable_name="history"),
            # HumanMessage(content="{result_prompt_text}"),
            HumanMessagePromptTemplate.from_template(template="{result_prompt_text}"),
        ]
    )
    # | log
)

chain = prompt_chain | create_llm() | StrOutputParser()

# 测试prompt_chain
# chain.invoke({"question": 123})
# chain.invoke({"question": "图像局部重绘的使用场景"})
# chain.invoke({"question": "图像局部重绘的Python代码实现"})


def get_input():
    return input("\n你好，请输入你的问题：\n")


input_text = get_input()

while input_text != "exit":
    # answer = chain.invoke({"question": input_text})
    # print(answer)
    full_text = ""
    for chunk in chain.stream({"question": input_text}):
        full_text = full_text + chunk
        print(chunk, end='', flush=True)
    input_text = get_input()
    save_history_item(HumanMessage(content=input_text))
    save_history_item(AIMessage(content=full_text))
